﻿using System;
using System.Collections.Generic;
using System.Text;

namespace MLForgeSharp.Models.ProbabilisticModels.PGM
{
    using System;

    /// <summary>
    /// 隐马尔科夫模型
    /// </summary>
    public class HiddenMarkovModel
    {
        private int numStates; // 状态数
        private int numObservations; // 观测数
        private double[] initialProbabilities; // 初始概率
        private double[,] transitionProbabilities; // 转移概率
        private double[,] emissionProbabilities; // 发射概率

        public HiddenMarkovModel(int numStates, int numObservations)
        {
            this.numStates = numStates;
            this.numObservations = numObservations;

            // 初始化概率矩阵
            initialProbabilities = new double[numStates];
            transitionProbabilities = new double[numStates, numStates];
            emissionProbabilities = new double[numStates, numObservations];

            Random rand = new Random();
            for (int i = 0; i < numStates; i++)
            {
                initialProbabilities[i] = 1.0 / numStates; // 均匀分布
                for (int j = 0; j < numStates; j++)
                {
                    transitionProbabilities[i, j] = 1.0 / numStates; // 均匀分布
                }
                for (int k = 0; k < numObservations; k++)
                {
                    emissionProbabilities[i, k] = 1.0 / numObservations; // 均匀分布
                }
            }
        }

        // Viterbi 算法：预测最可能的状态序列
        public int[] Predict(int[] observations)
        {
            int T = observations.Length;
            double[,] delta = new double[T, numStates];
            int[,] psi = new int[T, numStates];

            // 初始化
            for (int i = 0; i < numStates; i++)
            {
                delta[0, i] = initialProbabilities[i] * emissionProbabilities[i, observations[0]];
                psi[0, i] = 0;
            }

            // 递推
            for (int t = 1; t < T; t++)
            {
                for (int j = 0; j < numStates; j++)
                {
                    double maxProb = 0.0;
                    int maxState = 0;
                    for (int i = 0; i < numStates; i++)
                    {
                        double prob = delta[t - 1, i] * transitionProbabilities[i, j];
                        if (prob > maxProb)
                        {
                            maxProb = prob;
                            maxState = i;
                        }
                    }
                    delta[t, j] = maxProb * emissionProbabilities[j, observations[t]];
                    psi[t, j] = maxState;
                }
            }

            // 回溯
            int[] states = new int[T];
            double maxFinalProb = 0.0;
            for (int i = 0; i < numStates; i++)
            {
                if (delta[T - 1, i] > maxFinalProb)
                {
                    maxFinalProb = delta[T - 1, i];
                    states[T - 1] = i;
                }
            }
            for (int t = T - 2; t >= 0; t--)
            {
                states[t] = psi[t + 1, states[t + 1]];
            }

            return states;
        }

        // Baum-Welch 算法：训练模型
        public void Train(int[] observations, int maxIterations = 100)
        {
            int T = observations.Length;

            for (int iter = 0; iter < maxIterations; iter++)
            {
                double[,] alpha = new double[T, numStates];
                double[,] beta = new double[T, numStates];
                double[,] gamma = new double[T, numStates];
                double[, ,] xi = new double[T - 1, numStates, numStates];

                // 前向算法
                for (int i = 0; i < numStates; i++)
                {
                    alpha[0, i] = initialProbabilities[i] * emissionProbabilities[i, observations[0]];
                }
                for (int t = 1; t < T; t++)
                {
                    for (int j = 0; j < numStates; j++)
                    {
                        alpha[t, j] = 0.0;
                        for (int i = 0; i < numStates; i++)
                        {
                            alpha[t, j] += alpha[t - 1, i] * transitionProbabilities[i, j];
                        }
                        alpha[t, j] *= emissionProbabilities[j, observations[t]];
                    }
                }

                // 后向算法
                for (int i = 0; i < numStates; i++)
                {
                    beta[T - 1, i] = 1.0;
                }
                for (int t = T - 2; t >= 0; t--)
                {
                    for (int i = 0; i < numStates; i++)
                    {
                        beta[t, i] = 0.0;
                        for (int j = 0; j < numStates; j++)
                        {
                            beta[t, i] += beta[t + 1, j] * transitionProbabilities[i, j] * emissionProbabilities[j, observations[t + 1]];
                        }
                    }
                }

                // 计算 gamma 和 xi
                for (int t = 0; t < T - 1; t++)
                {
                    double sum = 0.0;
                    for (int i = 0; i < numStates; i++)
                    {
                        for (int j = 0; j < numStates; j++)
                        {
                            xi[t, i, j] = alpha[t, i] * transitionProbabilities[i, j] * emissionProbabilities[j, observations[t + 1]] * beta[t + 1, j];
                            sum += xi[t, i, j];
                        }
                    }
                    for (int i = 0; i < numStates; i++)
                    {
                        for (int j = 0; j < numStates; j++)
                        {
                            xi[t, i, j] /= sum;
                        }
                    }
                }
                for (int t = 0; t < T; t++)
                {
                    double sum = 0.0;
                    for (int i = 0; i < numStates; i++)
                    {
                        gamma[t, i] = alpha[t, i] * beta[t, i];
                        sum += gamma[t, i];
                    }
                    for (int i = 0; i < numStates; i++)
                    {
                        gamma[t, i] /= sum;
                    }
                }

                // 更新模型参数
                for (int i = 0; i < numStates; i++)
                {
                    initialProbabilities[i] = gamma[0, i];
                }
                for (int i = 0; i < numStates; i++)
                {
                    for (int j = 0; j < numStates; j++)
                    {
                        double numerator = 0.0;
                        double denominator = 0.0;
                        for (int t = 0; t < T - 1; t++)
                        {
                            numerator += xi[t, i, j];
                            denominator += gamma[t, i];
                        }
                        transitionProbabilities[i, j] = numerator / denominator;
                    }
                }
                for (int i = 0; i < numStates; i++)
                {
                    for (int k = 0; k < numObservations; k++)
                    {
                        double numerator = 0.0;
                        double denominator = 0.0;
                        for (int t = 0; t < T; t++)
                        {
                            if (observations[t] == k)
                            {
                                numerator += gamma[t, i];
                            }
                            denominator += gamma[t, i];
                        }
                        emissionProbabilities[i, k] = numerator / denominator;
                    }
                }
            }
        }
    }

    // 示例程序
    public class HiddenMarkovModelExample
    {
        public HiddenMarkovModelExample()
        {
            // 示例数据
            int[] observations = { 0, 1, 0, 1, 0 }; // 观测序列

            // 创建 HMM 模型
            HiddenMarkovModel hmm = new HiddenMarkovModel(numStates: 2, numObservations: 2);

            // 训练模型
            hmm.Train(observations);

            // 预测状态序列
            int[] predictedStates = hmm.Predict(observations);

            Console.WriteLine("Predicted States: ");
            foreach (int state in predictedStates)
            {
                Console.Write(state + ' ');
            }
        }
    }
}
